##########################################################################################

library(data.table)
library(optparse)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--peak_gene_file"), type = "character"),
    make_option(c("--exp_score_file"), type = "character"),
    make_option(c("--nclust"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    peak_gene_file <- "~/20231121_singleMuti/results/tmp_qc_v2/celltype_plot/peak2gene/germ/peakToGeneHeatmap_LabelClust_k25.peakNum.tsv"

    ## 表达和开放性的输入
    exp_score_file <- "~/20231121_singleMuti/results/tmp_qc_v2/celltype_plot/exp_score/germ_exp-atac_linkPeakNum.tsv"

    ## 
    nclust <- 11

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/tmp_qc_v2/celltype_plot/peak2gene/germ"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

peak_gene_file <- opt$peak_gene_file
exp_score_file <- opt$exp_score_file
scriptPath <- opt$scriptPath
nclust <- opt$nclust
out_path <- opt$out_path

dir.create( out_path , recursive = T)

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

##########################################################################################
## 读入文件
peak_gene <- data.frame(fread(peak_gene_file))
exp_score <- data.frame(fread(exp_score_file))

##########################################################################################
## 通路富集分析，提供总的基因以及关注的基因即可
goclust <- function(all_genes = all_genes , clust_genes = clust_genes , k = k){
  upGO <- rbind(
    calcTopGo(all_genes, interestingGenes=clust_genes, nodeSize=5, ontology="BP") 
    )

  upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

  ## 构造GO对象
  geneList <- factor(as.integer(all_genes %in% clust_genes))
  names(geneList) <- all_genes
  GOdata <- suppressMessages(new(
    "topGOdata",
    ontology = "BP",
    allGenes = geneList,
    annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
    nodeSize = 5
  ))
  ## 提取每个通路中感兴趣的基因
  gene_in_pathway <- sapply(upGO$GO.ID, function(x)
    {
      genes <- genesInTerm(GOdata, x)
      # myGenes is the queried gene list
      paste0(genes[[1]][genes[[1]] %in% clust_genes] , collapse = "," )
    })
  
  upGO$gene_in_pathway <- gene_in_pathway
  upGO$cluster <- paste0("cluster" , k)
  return(upGO)
}

##########################################################################################
## 分10%的peak数量的基因，以及gene其本身开放-表达显著正相关且存在至少3个peak连接的基因
data <- merge( peak_gene , exp_score , by = "gene" , all.x = T )
all_genes <- unique(data$gene)
## 前10%
nGOgenes <- 0.1
## 表达-开放的相关系数
cor_t <- 0.5
min_peak_num <- 3

##
result_peak10 <- c()
result_cor <- c()

for( k in unique(data$kclust) ){

  #### 前10%的peak数量最高的基因
  tmp <- subset( data , kclust == k )
  ## 按照peak的数量排序
  tmp <- tmp[order(tmp$peak_num , decreasing = T),]
  tmp <- tmp[1:round(nrow(tmp) * 0.1),]
  clust_genes <- tmp$gene
  ## 通路富集
  tmp_go <- goclust(all_genes = all_genes , clust_genes = clust_genes , k = k)
  tmp_go_peak10 <- tmp_go
  
  #### 取基因的表达和开放性存在相关的基因，且至少存在3个peak的连接
  clust_genes <- subset(tmp , cor_exp2atac >= cor_t & fdr_exp2atac < 0.05 & peak_num >= min_peak_num )$gene
  tmp_go <- goclust(all_genes = all_genes , clust_genes = clust_genes , k = k)
  tmp_go_cor <- tmp_go

  result_peak10 <- rbind( result_peak10 , tmp_go_peak10 )
  result_cor <- rbind( result_cor , tmp_go_cor )
}

out_file <- paste0(out_path, "/peakToGeneHeatmap_LabelClust_k" , nclust , ".peak_10.tsv")
write.table( result_peak10 , out_file , row.names = F , quote = F , sep ="\t" )
out_file <- paste0(out_path, "/peakToGeneHeatmap_LabelClust_k" , nclust , ".cor_exp2atacB0.5_peakB2.tsv")
write.table( result_cor , out_file , row.names = F , quote = F , sep ="\t" )


##########################################################################################
## 通路富集分析画图
# Plots of GO term enrichments:
goplot <- function(out_file = out_file , GOresults = GOresults){
  pdf(out_file , width=10, height=24)
  for(name in  paste0("cluster" , seq(1:25))){
      goRes <- subset( GOresults , cluster == name )
      if(nrow(goRes)>1){
        print(topGObarPlot(goRes, cmap = cmaps_BOR$comet, 
          nterms=40, border_color="black", 
          barwidth=0.85, title=name, barLimits=c(0, 15)))
      }
  }
  dev.off()
}

GOresults <- result_peak10
out_file <- paste0(out_path , "/peakToGeneHeatmap_LabelClust_k" , nclust , ".peak_10.pdf")
goplot(out_file = out_file , GOresults = GOresults)

GOresults <- result_cor
out_file <- paste0(out_path, "/peakToGeneHeatmap_LabelClust_k" , nclust , ".cor_exp2atacB0.5_peakB2.pdf")
goplot(out_file = out_file , GOresults = GOresults)
